package com.smy.baseutil.lock;

import java.io.IOException;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

import org.I0Itec.zkclient.IZkDataListener;
import org.I0Itec.zkclient.ZkClient;
import org.I0Itec.zkclient.exception.ZkNoNodeException;
import org.I0Itec.zkclient.serialize.SerializableSerializer;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.smy.baseutil.config.Property;


/**
 * 
 * @Title: ZkLock.java
 * @Description: zookeeper实现分布式锁
 * @author zhoutao
 * @date 2018年2月27日 上午10:26:18
 * @version V1.0
 */
public class ZkLock {

	private static Logger log = LoggerFactory.getLogger(RedisLock.class);

	public static final Long DEFAULT_TIMES = 10L;

	/**
	 * 锁根目录
	 */
	private static final String LOCK_ROOT_PATH = "/lock";

	/**
	 * 最大重试次数
	 */
	private static final Integer MAX_RETRY_COUNT = 10;

	/**
	 * ZK客户端
	 */
	private static final ZkClient ZK_CLIENT;

	/**
	 * 应用名
	 */
	private static final String APP_NAME;

	/**
	 * 用于保存Zookeeper中实现分布式锁的节点，如名称为locker：/locker，
	 * 该节点应该是持久节点，在该节点下面创建临时顺序节点来实现分布式锁
	 */
	private final String basePath;

	/**
	 * 用于保存某个客户端在locker下面创建成功的顺序节点，用于后续相关操作使用（如判断）
	 */
	private String ourLockPath;

	/**
	 * 锁名
	 */
	private String lockName;

	// 初始化，获取ZK地址及当前应用名
	static {
		String dubboZkAddress = Property.getProperty("dubbo.registry.address");
		APP_NAME = Property.getProperty("dubbo.application.name");
		if (StringUtils.isBlank(dubboZkAddress)) {
			throw new RuntimeException("无法获取:dubbo.registry.address的配置");
		}
		if (StringUtils.isBlank(APP_NAME)) {
			throw new RuntimeException("无法获取:dubbo.application.name的配置值。");
		}
		String zkAddress = dubboZkAddress.replaceAll("zookeeper://", "").replaceAll("\\?backup=", ",");
		ZK_CLIENT = new ZkClient(zkAddress, 10000, 10000, new SerializableSerializer());
	}

	/**
	 * 传入Zookeeper客户端连接对象，和basePath
	 *
	 * @param lockName
	 *            锁名
	 */
	public ZkLock(String lockName) {

		if (!ZK_CLIENT.exists(LOCK_ROOT_PATH)) {
			ZK_CLIENT.createPersistent(LOCK_ROOT_PATH);
		}
		this.basePath = LOCK_ROOT_PATH.concat("/").concat(APP_NAME);
		if (!ZK_CLIENT.exists(basePath)) {
			ZK_CLIENT.createPersistent(basePath);
		}
		this.ourLockPath = basePath.concat("/").concat(lockName);
		this.lockName = lockName;
	}

	/**
	 * 获取锁，直到超时，超时后抛出异常
	 *
	 * @throws Exception
	 *             异常
	 */
	public void lock() throws Exception {
		if (!internalLock(-1, null)) {
			throw new IOException("连接丢失!在路径:'" + basePath + "'下不能获取锁!");
		}

	}

	/**
	 * 获取锁，带有超时时间
	 *
	 * @param time
	 *            超时时长
	 * @param unit
	 *            时间单位
	 * @return true/false
	 */
	public boolean lock(long time, TimeUnit unit) {
		try {
			return internalLock(time, unit);
		} catch (Exception e) {
			log.error("获取锁异常。", e);
			return false;
		}
	}

	/**
	 * 释放锁
	 *
	 * @throws Exception
	 *             异常
	 */
	public void unlock() throws Exception {
		try {
			releaseLock(ourLockPath);
		} catch (Exception e) {
			log.error("释放锁异常。", e);
		}
	}

	/**
	 * 用于获取锁资源，通过父类的获取锁方法来获取锁
	 *
	 * @param time
	 *            获取锁的超时时间
	 * @param unit
	 *            time的时间单位
	 * @return 是否获取到锁
	 * @throws Exception
	 *             异常
	 */
	private boolean internalLock(long time, TimeUnit unit) throws Exception {
		// 如果ourLockPath不为空则认为获取到了锁，具体实现细节见attemptLock的实现
		ourLockPath = attemptLock(time, unit);
		return ourLockPath != null;
	}

	/**
	 * 尝试加锁
	 *
	 * @param time
	 *            超时时长
	 * @param unit
	 *            时间单位
	 * @return 锁名
	 * @throws Exception
	 *             异常
	 */
	private String attemptLock(long time, TimeUnit unit) throws Exception {
		final long startMillis = System.currentTimeMillis();
		final Long millisToWait = (unit != null) ? unit.toMillis(time) : null;

		String ourPath = null;
		boolean hasTheLock = false;
		boolean isDone = false;
		int retryCount = 0;

		// 网络闪断需要重试一试
		while (!isDone) {
			isDone = true;

			try {
				// createLockNode用于在locker（basePath持久节点）下创建客户端要获取锁的[临时]顺序节点
				ourPath = createLockNode(ourLockPath);

				// 该方法用于判断自己是否获取到了锁，即自己创建的顺序节点在locker的所有子节点中是否最小
				// 如果没有获取到锁，则等待其它客户端锁的释放，并且稍后重试直到获取到锁或者超时
				hasTheLock = waitToLock(startMillis, millisToWait, ourPath);

			} catch (ZkNoNodeException e) {
				if (retryCount++ < MAX_RETRY_COUNT) {
					isDone = false;
				} else {
					throw e;
				}
			}
		}
		if (hasTheLock) {
			return ourPath;
		}
		return null;
	}

	/**
	 * 释放锁
	 *
	 * @param lockPath
	 *            锁路径
	 * @throws Exception
	 *             异常
	 */
	private void releaseLock(String lockPath) throws Exception {
		ZK_CLIENT.delete(lockPath);
	}

	/**
	 * 创建锁节点
	 *
	 * @param path
	 *            路径
	 * @return 路径名
	 * @throws Exception
	 *             异常
	 */
	private String createLockNode(String path) throws Exception {
		return ZK_CLIENT.createEphemeralSequential(path, null);
	}

	/**
	 * 加锁
	 *
	 * @param startMillis
	 *            开始时间
	 * @param millisToWait
	 *            超时时间
	 * @param ourPath
	 *            路径
	 * @return true/false
	 * @throws Exception
	 *             异常
	 */
	private boolean waitToLock(long startMillis, Long millisToWait, String ourPath) throws Exception {
		boolean haveTheLock = false;
		boolean doDelete = false;

		try {
			while (!haveTheLock) {

				// 该方法实现获取locker节点下的所有顺序节点，并且从小到大排序
				List<String> children = getSortedChildren();
				String sequenceNodeName = ourPath.substring(basePath.length() + 1);

				// 计算刚才客户端创建的顺序节点在locker的所有子节点中排序位置，如果是排序为0，则表示获取到了锁
				int ourIndex = children.indexOf(sequenceNodeName);

				// 如果在getSortedChildren中没有找到之前创建的[临时]顺序节点，这表示可能由于网络闪断而导致
				// Zookeeper认为连接断开而删除了我们创建的节点，此时需要抛出异常，让上一级去处理
				// 上一级的做法是捕获该异常，并且执行重试指定的次数 见后面的 attemptLock方法
				if (ourIndex < 0) {
					throw new ZkNoNodeException("节点没有找到: " + sequenceNodeName);
				}

				// 如果当前客户端创建的节点在locker子节点列表中位置大于0，表示其它客户端已经获取了锁,此时当前客户端需要等待其它客户端释放锁.
				boolean isGetTheLock = ourIndex == 0;

				// 如何判断其它客户端是否已经释放了锁？从子节点列表中获取到比自己次小的哪个节点，并对其建立监听
				String pathToWatch = isGetTheLock ? null : children.get(ourIndex - 1);

				if (isGetTheLock) {
					haveTheLock = true;
				} else {

					// 如果次小的节点被删除了，则表示当前客户端的节点应该是最小的了，所以使用CountDownLatch来实现等待
					String previousSequencePath = basePath.concat("/").concat(pathToWatch);
					final CountDownLatch latch = new CountDownLatch(1);
					final IZkDataListener previousListener = new IZkDataListener() {

						/**
						 * 次小节点删除事件发生时，让countDownLatch结束等待.此时还需要重新让程序回到while，
						 * 重新判断一次！
						 * 
						 * @param dataPath
						 *            数据路径
						 * @throws Exception
						 *             异常
						 */
						@Override
						public void handleDataDeleted(String dataPath) throws Exception {
							latch.countDown();
						}

						@Override
						public void handleDataChange(String dataPath, Object data) throws Exception {
							// ignore
						}
					};

					try {
						// 如果节点不存在会出现异常
						ZK_CLIENT.subscribeDataChanges(previousSequencePath, previousListener);
						if (millisToWait != null) {
							millisToWait -= (System.currentTimeMillis() - startMillis);
							startMillis = System.currentTimeMillis();
							if (millisToWait <= 0) {
								// timed out - delete our node
								doDelete = true;
								break;
							}
							latch.await(millisToWait, TimeUnit.MICROSECONDS);
						} else {
							latch.await();
						}
					} catch (ZkNoNodeException e) {
						// ignore
					} finally {
						ZK_CLIENT.unsubscribeDataChanges(previousSequencePath, previousListener);
					}
				}
			}
		} catch (Exception e) {
			// 发生异常需要删除节点
			doDelete = true;
			throw e;
		} finally {
			// 如果需要删除节点
			if (doDelete) {
				releaseLock(ourPath);
			}
		}
		return haveTheLock;
	}

	/**
	 * 获取锁节点Sequence编号
	 *
	 * @param str
	 *            路径
	 * @param lockName
	 *            锁名称
	 * @return 锁Sequence
	 */
	private String getLockNodeNumber(String str, String lockName) {
		int index = str.lastIndexOf(lockName);
		if (index >= 0) {
			index += lockName.length();
			return index <= str.length() ? str.substring(index) : "";
		}
		return str;
	}

	/**
	 * 子节点排序
	 *
	 * @return 排序后的节点列表
	 * @throws Exception
	 *             异常
	 */
	private List<String> getSortedChildren() throws Exception {
		try {
			List<String> children = ZK_CLIENT.getChildren(basePath);
			Collections.sort(children, new Comparator<String>() {

				@Override
				public int compare(String lhs, String rhs) {
					return getLockNodeNumber(lhs, lockName).compareTo(getLockNodeNumber(rhs, lockName));
				}
			});
			return children;
		} catch (ZkNoNodeException e) {
			ZK_CLIENT.createPersistent(basePath, true);
			return getSortedChildren();
		}
	}

}
